package com.riteny.oauth.core.tokenstore;

import com.riteny.oauth.core.userauth.CustomUser;
import com.riteny.oauth.core.userauth.CustomUserAuthenticationService;
import org.springframework.dao.DuplicateKeyException;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.support.SqlLobValue;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.DefaultAuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.store.JdbcTokenStore;

import javax.sql.DataSource;
import java.sql.Types;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author Riteny
 * 2021/5/7  17:39
 */
public class CustomJdbcTokenStore extends JdbcTokenStore implements CustomTokenStore {

    private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();

    private static final String DEFAULT_ACCESS_TOKEN_FROM_AUTHENTICATION_SELECT_STATEMENT = "select token_id, token from oauth_access_token where authentication_id = ?";

    private static final String DEFAULT_ACCESS_TOKEN_INSERT_STATEMENT = "insert into oauth_access_token (token_id, token, authentication_id, user_name, client_id, authentication, refresh_token) values (?, ?, ?, ?, ?, ?, ?)";

    private final JdbcTemplate jdbcTemplate;

    private CustomUserAuthenticationService customUserAuthenticationService;

    public CustomJdbcTokenStore(DataSource dataSource, CustomUserAuthenticationService customUserAuthenticationService) {
        super(dataSource);
        this.customUserAuthenticationService = customUserAuthenticationService;
        this.jdbcTemplate = new JdbcTemplate(dataSource);
    }

    @Override
    public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {

        String refreshToken = null;
        if (token.getRefreshToken() != null) {
            refreshToken = token.getRefreshToken().getValue();
        }

        try {
            jdbcTemplate.update(DEFAULT_ACCESS_TOKEN_INSERT_STATEMENT, new Object[]{extractTokenKey(token.getValue()),
                    new SqlLobValue(serializeAccessToken(token)), authenticationKeyGenerator.extractKey(authentication),
                    authentication.isClientOnly() ? null : authentication.getName(),
                    authentication.getOAuth2Request().getClientId(),
                    new SqlLobValue(serializeAuthentication(authentication)), extractTokenKey(refreshToken)}, new int[]{
                    Types.VARCHAR, Types.BLOB, Types.VARCHAR, Types.VARCHAR, Types.VARCHAR, Types.BLOB, Types.VARCHAR});

        } catch (DuplicateKeyException e) {
            //主鍵重複，代表數據庫已經存在一個有效的token，此問題可能由於統一用戶同時間多綫程調用獲取token所致
            //從數據庫獲取已經存在的有效token，並覆蓋當前OAuth2AccessToken的内容
            replaceOauthTokenFromDatabase(token, authentication);
        }
    }

    private void replaceOauthTokenFromDatabase(OAuth2AccessToken token, OAuth2Authentication authentication) {

        DefaultOAuth2AccessToken oAuth2AccessToken = (DefaultOAuth2AccessToken) token;

        OAuth2AccessToken accessToken = jdbcTemplate.queryForObject(DEFAULT_ACCESS_TOKEN_FROM_AUTHENTICATION_SELECT_STATEMENT,
                (rs, rowNum) -> deserializeAccessToken(rs.getBytes(2)), authenticationKeyGenerator.extractKey(authentication));

        if (accessToken != null) {
            oAuth2AccessToken.setValue(accessToken.getValue());
            oAuth2AccessToken.setAdditionalInformation(accessToken.getAdditionalInformation());
            oAuth2AccessToken.setExpiration(accessToken.getExpiration());
            oAuth2AccessToken.setRefreshToken(accessToken.getRefreshToken());
            oAuth2AccessToken.setScope(accessToken.getScope());
            oAuth2AccessToken.setTokenType(accessToken.getTokenType());
        }
    }

    @Override
    public Map<String, Object> readUserInfo(OAuth2AccessToken token, OAuth2Authentication authentication) {
        CustomUser map = (CustomUser) authentication.getUserAuthentication().getPrincipal();
        String userType = map.getUserType();
        String username = map.getUsername();

        Map<String, Object> userInfo = new HashMap<>();
        userInfo.put("userDetails", customUserAuthenticationService.searchUserDetailInfo(username, userType));

        return userInfo;
    }

    public OAuth2AccessToken readAccessToken(String tokenValue) {
        return super.readAccessToken(tokenValue);
    }
}
